#!/usr/bin/env python3
"""
Percona XtraBackup MCP Server

面向 percona-xtrabackup 软件的 MCP 服务器，提供数据库备份、恢复、流处理和加密解密等命令行操作的封装。

MCP tools 列表：
1. xtrabackup_backup - 执行数据库备份
   - 参数：target_dir, compress, encrypt, encrypt_key, parallel, tmpdir, incremental_lsn, innodb_buffer_pool_size
   - 返回：执行结果 JSON

2. xtrabackup_prepare - 准备备份数据
   - 参数：target_dir, apply_log_only, export, rebuild_indexes, rebuild_threads, innodb_checksum_algorithm, innodb_log_checksum_algorithm, innodb_undo_directory, innodb_undo_tablespaces
   - 返回：执行结果 JSON

3. xbstream_create - 创建 xbstream 归档
   - 参数：files, output_file
   - 返回：执行结果 JSON

4. xbstream_extract - 从 xbstream 归档提取文件
   - 参数：input_file, verbose
   - 返回：执行结果 JSON

5. xbcrypt_encrypt - 加密文件
   - 参数：input_file, output_file, algorithm, key
   - 返回：执行结果 JSON

6. xbcrypt_decrypt - 解密文件
   - 参数：input_file, output_file, algorithm, key
   - 返回：执行结果 JSON

7. innobackupex_decrypt - 解密备份文件
   - 参数：backup_dir, decrypt_algorithm, encrypt_key_file
   - 返回：执行结果 JSON
"""

import subprocess
import json
from typing import List, Optional
from mcp.server.fastmcp import FastMCP, Context

mcp = FastMCP("Percona XtraBackup MCP Server")

@mcp.tool()
def xtrabackup_backup(
    target_dir: str,
    compress: bool = False,
    encrypt: Optional[str] = None,
    encrypt_key: Optional[str] = None,
    parallel: Optional[int] = None,
    tmpdir: Optional[str] = None,
    incremental_lsn: Optional[str] = None,
    innodb_buffer_pool_size: Optional[str] = None
) -> dict:
    """
    执行数据库备份操作
    
    参数:
        target_dir: 备份目标目录
        compress: 是否启用压缩
        encrypt: 加密算法（如 AES256）
        encrypt_key: 加密密钥
        parallel: 并行线程数
        tmpdir: 临时目录路径
        incremental_lsn: 增量备份的 LSN
        innodb_buffer_pool_size: InnoDB 缓冲池大小
    
    返回:
        {"success": bool, "command": str, "exit_code": int, "stdout": str, "stderr": str}
    """
    cmd = ["xtrabackup", "--backup", f"--target-dir={target_dir}"]
    
    if compress:
        cmd.append("--compress")
    if encrypt:
        cmd.extend(["--encrypt", encrypt])
    if encrypt_key:
        cmd.extend(["--encrypt-key", encrypt_key])
    if parallel:
        cmd.extend(["--parallel", str(parallel)])
    if tmpdir:
        cmd.extend(["--tmpdir", tmpdir])
    if incremental_lsn:
        cmd.extend(["--incremental-lsn", incremental_lsn])
    if innodb_buffer_pool_size:
        cmd.extend(["--innodb-buffer-pool-size", innodb_buffer_pool_size])
    
    command_str = " ".join(cmd)
    result = subprocess.run(cmd, capture_output=True, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": command_str,
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def xtrabackup_prepare(
    target_dir: str,
    apply_log_only: bool = False,
    export: bool = False,
    rebuild_indexes: bool = False,
    rebuild_threads: Optional[int] = None,
    innodb_checksum_algorithm: Optional[str] = None,
    innodb_log_checksum_algorithm: Optional[str] = None,
    innodb_undo_directory: Optional[str] = None,
    innodb_undo_tablespaces: Optional[int] = None
) -> dict:
    """
    准备备份数据
    
    参数:
        target_dir: 备份目录路径
        apply_log_only: 仅应用日志
        export: 导出表空间
        rebuild_indexes: 重建索引
        rebuild_threads: 重建线程数
        innodb_checksum_algorithm: InnoDB 校验和算法
        innodb_log_checksum_algorithm: InnoDB 日志校验和算法
        innodb_undo_directory: InnoDB undo 目录
        innodb_undo_tablespaces: InnoDB undo 表空间数量
    
    返回:
        {"success": bool, "command": str, "exit_code": int, "stdout": str, "stderr": str}
    """
    cmd = ["xtrabackup", "--prepare", f"--target-dir={target_dir}"]
    
    if apply_log_only:
        cmd.append("--apply-log-only")
    if export:
        cmd.append("--export")
    if rebuild_indexes:
        cmd.append("--rebuild-indexes")
    if rebuild_threads:
        cmd.extend(["--rebuild-threads", str(rebuild_threads)])
    if innodb_checksum_algorithm:
        cmd.extend(["--innodb-checksum-algorithm", innodb_checksum_algorithm])
    if innodb_log_checksum_algorithm:
        cmd.extend(["--innodb-log-checksum-algorithm", innodb_log_checksum_algorithm])
    if innodb_undo_directory:
        cmd.extend(["--innodb-undo-directory", innodb_undo_directory])
    if innodb_undo_tablespaces:
        cmd.extend(["--innodb-undo-tablespaces", str(innodb_undo_tablespaces)])
    
    command_str = " ".join(cmd)
    result = subprocess.run(cmd, capture_output=True, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": command_str,
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def xbstream_create(files: List[str], output_file: str) -> dict:
    """
    创建 xbstream 归档文件
    
    参数:
        files: 要包含在归档中的文件列表
        output_file: 输出文件路径
    
    返回:
        {"success": bool, "command": str, "exit_code": int, "stdout": str, "stderr": str}
    """
    cmd = ["xbstream", "-c"] + files
    
    command_str = f"{' '.join(cmd)} > {output_file}"
    with open(output_file, 'w') as f:
        result = subprocess.run(cmd, stdout=f, stderr=subprocess.PIPE, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": command_str,
        "exit_code": result.returncode,
        "stdout": "",
        "stderr": result.stderr
    }

@mcp.tool()
def xbstream_extract(input_file: str, verbose: bool = False) -> dict:
    """
    从 xbstream 归档中提取文件
    
    参数:
        input_file: 输入归档文件路径
        verbose: 是否显示详细输出
    
    返回:
        {"success": bool, "command": str, "exit_code": int, "stdout": str, "stderr": str}
    """
    cmd = ["xbstream", "-x"]
    if verbose:
        cmd.append("-v")
    
    command_str = f"{' '.join(cmd)} < {input_file}"
    with open(input_file, 'r') as f:
        result = subprocess.run(cmd, stdin=f, capture_output=True, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": command_str,
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def xbcrypt_encrypt(input_file: str, output_file: str, algorithm: str, key: str) -> dict:
    """
    加密文件
    
    参数:
        input_file: 输入文件路径
        output_file: 输出文件路径
        algorithm: 加密算法（如 aes-256-cbc）
        key: 加密密钥
    
    返回:
        {"success": bool, "command": str, "exit_code": int, "stdout": str, "stderr": str}
    """
    cmd = ["xbcrypt", "-i", input_file, "-o", output_file, "-a", algorithm, "-k", key]
    
    command_str = " ".join(cmd)
    result = subprocess.run(cmd, capture_output=True, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": command_str,
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def xbcrypt_decrypt(input_file: str, output_file: str, algorithm: str, key: str) -> dict:
    """
    解密文件
    
    参数:
        input_file: 输入文件路径
        output_file: 输出文件路径
        algorithm: 加密算法（如 aes-256-cbc）
        key: 解密密钥
    
    返回:
        {"success": bool, "command": str, "exit_code": int, "stdout": str, "stderr": str}
    """
    cmd = ["xbcrypt", "-d", "-i", input_file, "-o", output_file, "-a", algorithm, "-k", key]
    
    command_str = " ".join(cmd)
    result = subprocess.run(cmd, capture_output=True, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": command_str,
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def innobackupex_decrypt(backup_dir: str, decrypt_algorithm: str, encrypt_key_file: str) -> dict:
    """
    解密备份文件
    
    参数:
        backup_dir: 备份目录路径
        decrypt_algorithm: 解密算法（aes/twofish/camellia）
        encrypt_key_file: 加密密钥文件路径
    
    返回:
        {"success": bool, "command": str, "exit_code": int, "stdout": str, "stderr": str}
    """
    cmd = ["innobackupex", f"--decrypt={decrypt_algorithm}", f"--encrypt-key-file={encrypt_key_file}", backup_dir]
    
    command_str = " ".join(cmd)
    result = subprocess.run(cmd, capture_output=True, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": command_str,
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

if __name__ == "__main__":
    mcp.run()